# ---encoding:utf-8---
# @Time    : 2024/4/18 22:19
# @Author  : DuJingze
# @Email   ：
# @Site    : 
# @File    : ocr_template_match.py
# @Project : 12_GuPaoPytorch
# @Software: PyCharm
# from imutils import contours
import numpy as np
import argparse
import imutils
import cv2
import myutils

# 指定信用卡类型
FIRST_NUMBER = {
    "3": "American Express",
    "4": "Visa",
    "5": "MasterCard",
    "6": "Discover Card"
}


def zh_ch(string):
    return string.encode("gbk").decode(errors="ignore")


# 绘图展示
def cv_show(name, img):
    """
    展示cv类型图片数据
    :param name: 图片名称
    :param img: cv类型图片数据
    :return:
    """
    cv2.imshow(name, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


# 处理模板
def temp_deal(template_img_path):
    """
    处理模板，返回识别结果
    :param template_img_path: 模板图片文件存放路径
    :return:
    """
    # 1.读取模板图像
    temp_image = cv2.imread(template_img_path)
    cv_show("original_template_image", temp_image)
    # 2.将模板图像转换为灰度图
    temp_gray = cv2.cvtColor(temp_image, cv2.COLOR_BGR2GRAY)
    cv_show("template_grayscale_image", temp_gray)
    # 3.二值化处理，图像阈值函数，像素值超过127变成0，否则变成255  反二值化，高亮突出数字
    ret, temp_binary = cv2.threshold(temp_gray, 127, 255, cv2.THRESH_BINARY_INV)
    cv_show("template_image_binary", temp_binary)

    # 3.计算轮廓，返回轮廓信息和轮廓层数
    # 3.1 轮廓检测，参数:(二值化图像, 只检索最外层轮廓, (压缩水平、垂直和斜线，仅保留这些线段的端点))
    contours, hierarchy = cv2.findContours(temp_binary.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 3.2 绘制轮廓
    # 3.2.1 复制一份原图像作为画板，不能在原图上画，不然原图会改变
    draw_temp = temp_image.copy()
    # 3.2.2 "-1"是指在画板上画出所有轮廓信息，红色，线宽为2
    res = cv2.drawContours(draw_temp, contours, -1, (0, 0, 255), 2)
    cv_show("draw_all_outer_contours", res)

    # 4.模板轮廓排序
    sorted_bounding_list = myutils.sort_contours(contours, method="left-to-right")[0]  # 获取排序后的轮廓

    # 5.遍历每一个轮廓
    digits = {}
    for i, c in enumerate(sorted_bounding_list):  # 轮廓位置索引 及 对应的轮廓值
        # 5.1 计算外接矩形
        x, y, w, h = cv2.boundingRect(c)
        # 5.2 外接矩形对应的二值图像
        roi = temp_binary[y:y + h, x:x + w]
        # 5.3 roi区域统一大小，根据自己需求来定
        roi = cv2.resize(roi, (55, 85))
        # 5.4 显示当前矩形框区域图像
        window_name = "current_template:{}".format(i)
        cv_show(window_name, roi)
        # 5.5 保存当前数字的模板，每一个数字对应一个模板
        digits[i] = roi
    return digits


def deal_bank_card(digits, card_img_path):
    """
    处理待识别银行卡图片
    :param digits: 处理好的数字模板文件
    :param card_img_path: 银行卡图片路径
    :return:
    """
    # 初始化卷积核
    rectKernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 3))
    myKernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))

    img = cv2.imread("3.jpg")
    # cv_show("img", img)
    img = imutils.resize(img, width=300)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 礼帽操作，突出更明亮的区域
    tophat = cv2.morphologyEx(gray, cv2.MORPH_TOPHAT, rectKernel)
    # cv_show("tophat", tophat)

    # 利用Sobel算子进行边缘检测
    gradX = cv2.Sobel(tophat, ddepth=cv2.CV_32F, dx=1, dy=0, ksize=-1)
    gradY = cv2.Sobel(tophat, ddepth=cv2.CV_32F, dx=0, dy=1, ksize=-1)
    gradX = np.absolute(gradX)
    minVal = np.min(gradX)
    maxVal = np.max(gradX)
    gradX = (255 * ((gradX - minVal) / (maxVal - minVal)))
    gradX = gradX.astype("uint8")
    print(np.array(gradX).shape)
    # cv_show("gradX", gradX)

    # 通过闭操作，先膨胀后腐蚀，将数字连接在一块
    gradX = cv2.morphologyEx(gradX, cv2.MORPH_CLOSE, rectKernel)
    # cv_show("gradXclose", gradX)

    thresh = cv2.threshold(gradX, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
    # cv_show('threshImg', thresh)
    # 再来一个闭操作
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, myKernel)
    # cv_show('gradXclose2', thresh)

    # 计算轮廓
    threshCnts, hierarchy = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cnts = threshCnts
    cur_img = img.copy()
    cv2.drawContours(cur_img, cnts, -1, (0, 0, 255), 2)
    # cv_show('imgCopy', cur_img)

    locs = []

    # 遍历轮廓，提取轮廓，需要尝试
    for (i, c) in enumerate(cnts):
        # 计算矩形
        (x, y, w, h) = cv2.boundingRect(c)
        ar = w / float(h)
        # 选择合适的区域，根据实际任务来，这里的基本都是四个数字一组
        # [(454, 213, 110, 29), (308, 211, 109, 31), (171, 211, 98, 31), (18, 211, 107, 31)]
        if 2.5 < ar < 5.0:
            if (40 < w < 65) and (10 < h < 20):
                # 符合的留下来
                locs.append((x, y, w + 2, h + 2))
                res = cv2.rectangle(img.copy(), (x, y), (x + w + 2, y + h + 2), (0, 0, 255),
                                    2)  # "-1"是指在画板上画出所有轮廓信息，红色，线宽为2
                cv_show("imgRect", res)

    # 将符合的轮廓从左到右排序
    locs = sorted(locs, key=lambda o: o[0])
    output = []

    # 四组轮廓
    for (i, (gx, gy, gw, gh)) in enumerate(locs):
        groupOutput = []
        group = gray[gy - 5:gy + gh + 5, gx - 5:gx + gw + 5]
        group = cv2.threshold(group, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
        cv_show("groupThreshold", group)  # 1234；5678；
        digitCnts, his = cv2.findContours(group.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 提取轮廓内数字1,2,3,4
        digitCnts = myutils.sort_contours(digitCnts, method="left-to-right")[0]  # 对找到的轮廓进行排序  按照X轴坐标
        # 计算每一组中的每一个数值
        for c in digitCnts:  # 每组数字1,2,3,4
            # 找到当前数值的轮廓，resize成合适的大小
            (x, y, w, h) = cv2.boundingRect(c)  # 具体数字
            roi = group[y:y + h, x:x + w]  # 数字在每组group中的坐标
            # 将每个区域对应一个数字
            roi = cv2.resize(roi, (65, 99))  # roi区域统一大小，根据自己需求来定
            # cv_show('roi',roi)
            scores = []
            for (digitKey, digitROIValue) in digits.items():  # 识别图的roi数字与模板里每个数字matchTemplate
                result = cv2.matchTemplate(roi, digitROIValue, cv2.TM_CCOEFF_NORMED)
                # 返回最值及最值位置，在这里我们需要的是最小值的得分，不同的匹配度计算方法的选择不同
                min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
                scores.append(max_val)  # 最大值
            score = np.abs(scores)  # 有负数出现，统一成正数，相关系数都变成正数
            flag = False
            for val in score:
                if val > 0.9:
                    flag = True
                    break
            if not flag:
                best_index = np.argmax(score)  # score最大值的下标，匹配度最高
                best_value = str(best_index)  # 下标就是对应的数字，在字典中，key是0对应的是值为0的图片
                groupOutput.append(best_value)
        cv2.rectangle(img, (gx - 5, gy - 5), (gx + gw + 5, gy + gh + 5), (0, 0, 255), 1)
        # 在矩形框上绘图
        cv2.putText(img, ''.join(groupOutput), (gx, gy - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 255), 2)
        output.append(groupOutput)
    cv_show('img', img)
    # print('数字为：', output)
    for i in output:
        for j in i:
            print(j, end="")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--image_path", required=True, help="path of input image")
    parser.add_argument("-t", "--template_path", required=True, help="path of template OCR-A image")
    args = parser.parse_args()

    # 1.处理模板
    res_digits = temp_deal(args.template_path)

    # 2.处理银行卡图片
    deal_bank_card(res_digits, args.image_path)
